# @Time : 2024-11-07 16:19
# @Author : Fioman 
# @Phone : 13149920693
"""
图像处理过程中的帮助类,比如创建图像目录,显示图像之类的.
获取所有的图像啊之类的.
"""
import os
from collections import namedtuple

import cv2
import cv2 as cv
import imutils
import numpy as np

from public_tool.data_result import DataResult

fileDir = r"F:\raw_2025\4pin_newline"
keepTotalBoardResDir = None
keepOkDir = None
keepOkResDir = None
showState = True
sizeScale = 4
sensorPoint = (1639, 180)
sensorDisThres = 15
colorWhite = (255, 255, 255)
colorGreen = (0, 255, 0)
colorBlue = (255, 0, 0)
colorRed = (0, 0, 255)
colorPurple = (128, 0, 128)  # 紫色
mm2pix = 2.19923
colorList = [colorRed, colorGreen, colorBlue, colorPurple]


def show_image(name, image):
    if showState:
        h, w = image.shape[:2]
        newH, newW = h // sizeScale, w // sizeScale
        showImage = cv.resize(image, (newW, newH), interpolation=cv.INTER_AREA)
        cv.namedWindow(str(name))
        cv.imshow(str(name), showImage)
        cv.waitKey(0)


def create_dir_and_get_filenames(rootPath) -> list[str]:
    """
    创建图像目录,还有获取图像集合.
    :param rootPath:
    :return:
    """
    global keepTotalBoardResDir, keepOkDir, keepOkResDir, fileDir
    if not os.path.isdir(rootPath):
        fileDir, _ = os.path.split(rootPath)
    else:
        fileDir = rootPath
    keepTotalBoardResDir = os.path.join(fileDir, "total_board_res")
    keepOkDir = os.path.join(fileDir, "ok")
    keepOkResDir = os.path.join(fileDir, "ok_res")
    keepFailPath = os.path.join(fileDir, "fail")
    for dirName in [keepTotalBoardResDir, keepOkDir, keepOkResDir, keepFailPath]:
        if not os.path.isdir(dirName):
            os.makedirs(dirName)
    fileNames = []
    if os.path.isfile(rootPath):
        _, fileName = os.path.split(rootPath)
        fileNames.append(fileName)
    else:
        for root, dirs, files in os.walk(fileDir):
            if root != fileDir:
                continue

            for file in files:
                if not file.endswith(".bmp"):
                    print(f"文件名: {file}  错误,跳过")
                    continue
                fileNames.append(file)
    return fileNames


def get_center_by_contour(contour):
    """
    根据轮廓获取轮廓中心点,返回int类型的轮廓中心
    :param contour:
    :return:
    """
    (centerX, centerY), *_ = cv.minAreaRect(contour)
    return int(centerX + 0.5), int(centerY + 0.5)


def get_rect_points_clockwise(rectPoints):
    """
    获取四个点的顺时针方向的坐标,从左上角开始,到左下角
    :param rectPoints:
    :return:
    """
    rectPointsSorted = sorted(rectPoints, key=np.sum)
    leftTop = rectPointsSorted[0]
    rightBottom = rectPointsSorted[-1]
    secondPointX = rectPointsSorted[1][0]
    thirdPointX = rectPointsSorted[2][0]
    if secondPointX > thirdPointX:
        rightTop = rectPointsSorted[1]
        leftBottom = rectPointsSorted[2]
    else:
        rightTop = rectPointsSorted[2]
        leftBottom = rectPointsSorted[1]
    return leftTop, rightTop, rightBottom, leftBottom


def get_two_point_distance(p1, p2):
    """
    获取两个点之间的距离
    :param p1:
    :param p2:
    :return:
    """
    assert len(p1) == 2
    assert len(p2) == 2
    ds = np.sqrt((p2[0] - p1[0]) ** 2 + (p2[1] - p1[1]) ** 2)
    return round(ds, 2)


def get_two_lines_intersection(k1, b1, k2, b2):
    """
    获取两条直线的交点
    :param k1: 第一条直线的斜率
    :param b1: 第一条直线的斜距
    :param k2: 第二条直线的斜率
    :param b2: 第二条直线的斜距
    :return: 返回两条直线的交点
    """
    if k1 == k2:
        return None, None
    elif k1 is None and k2 is not None:
        x = b1
        y = k2 * x + b2
        return x, y
    elif k2 is None and k1 is not None:
        x = b2
        y = k1 * x + b1
        return x, y
    else:
        x = (b1 - b2) / (k1 - k2) * (-1)
        y = k1 * x + b1
        return x, y


def get_line_k_and_b(p1, p2):
    """
    根据两个点获取这条线的斜率和斜距
    :param p1: 线段的点1
    :param p2: 线段的点2
    :return:
    """
    if p1[0] == p2[0]:
        # 垂直于x轴,没有k值,斜率是None
        return None, p1[0]
    else:
        k = (p2[1] - p1[1]) / (p2[0] - p1[0])
        b = p1[1] - k * p1[0]
        return k, b


def get_point_to_line_distance(p, lineK, lineB):
    """
    获取点到直线的距离
    :param p: 要获取到直线的距离的点
    :param lineK: 直线的斜率K
    :param lineB: 直线的斜距B
    :return:
    """
    if lineK == 0:
        # 如果直线是水平的,那么它的垂线是垂直x轴的,也就是没有斜率.
        verLineK = None
        verLineB = p[0]
    elif lineK == None:
        verLineK = 0
        verLineB = p[1]
    else:
        # 垂线和原来的直线的斜率的乘积是-1
        verLineK = -1 / lineK
        verLineB = p[1] - verLineK * p[0]

    # 然后根据垂线的斜率和斜距以及原来的线的斜率和斜距来获取交点
    intersection = get_two_lines_intersection(lineK, lineB, verLineK, verLineB)
    # 然后是获取两个点的距离,距离有正负,如果检测的直线,在传感器点的下方为正,在上方为负
    dis = get_two_point_distance(p, intersection)
    dis = dis if p[1] > intersection[1] else -dis
    return round(dis, 2)


def get_two_points_middle(p1, p2):
    """
    获取两个点的中点坐标
    :param p1: 点1的坐标
    :param p2: 点2的坐标
    :return:
    """
    assert len(p1) == 2
    assert len(p2) == 2
    middleX = (p1[0] + p2[0]) / 2
    middleY = (p1[1] + p2[1]) / 2
    return round(middleX, 2), round(middleY, 2)


def get_board_without_angle(image, boxMaster, extend=50):
    """
    获取无角度的物料区域,是否延长,表示是否延长物料区域的四周,保证可以完整的看到物料区域.
    extend表示延长多少的意思.
    :param image:
    :param boxMaster:
    :param extend:延长多少的标识,0表示不延长
    :return:
    """
    (centerX, centerY), (w, h), angle = boxMaster
    widthHalf = int(w / 2)
    heightHalf = int(h / 2)
    # 之所以有零的情况,是因为旋转之后物料框的y起点可能跑到外面去了.
    yStart = max(int(centerY - heightHalf), 0)
    yEnd = min(int(centerY + heightHalf), image.shape[0])  # yEnd也是一样的,有可能出去了
    xStart = max(int(centerX - widthHalf), 0)
    xEnd = min(int(centerX + widthHalf), image.shape[1])
    M = cv.getRotationMatrix2D((centerX, centerY), -angle, 1)
    if extend != 0:
        if yStart - extend > 0:
            yStart -= extend
        if yEnd + extend < image.shape[0]:
            yEnd += extend

        if xStart - extend > 0:
            xStart -= extend
        if xStart + extend < image.shape[1]:
            xEnd += extend

    totalBoard = cv.warpAffine(image, M,
                               (image.shape[1], image.shape[0]))[yStart:yEnd, xStart:xEnd]
    return totalBoard, [xStart, yStart]


def check_totalboard_is_ok(image) -> DataResult:
    """
    检测物料区域是否符合要求
    :param image:
    :return:
    """
    imageColor = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    CheckData = namedtuple("CheckData", ["box", "offset", "boardRes",
                                         "boardResNoAngle", "totalboard"])
    result = DataResult(CheckData, None, None, None, None, None)
    thresROI = image[:170, :]
    show_image("ThresROI", thresROI)
    thresUsedVal = int(np.mean(thresROI) + 18)
    print(f"物料区域查找的阈值: {thresUsedVal}")
    filterArea = 4500
    T, board = cv.threshold(image, thresUsedVal, 255, cv.THRESH_BINARY)
    show_image(f"BoardThres_({thresUsedVal})", board)
    # 使用开操作,去除掉小点,开操作,先腐蚀再膨胀
    kernelSize = 3
    morphTimes = 5
    kernel = cv.getStructuringElement(cv.MORPH_RECT, (kernelSize, kernelSize))
    board = cv.morphologyEx(board, cv.MORPH_OPEN, kernel, iterations=morphTimes)
    show_image("BoardOpened", board)
    """
    cv.RETR_EXTERNAL: 只检索最外层的轮廓(即外部轮廓).这种模式会忽略图像中的内部轮廓(即嵌套的轮廓)
    cv.CHAIN_APPROX_SIMPLE:这种方法会压缩水平,垂直和对角线段,只保留端点.它将较长的直线段简化为
    两个点,而不是列出所有的轮廓点.
    速度: 这种方法通常更快,因为它减少了轮廓中的点数,降低了计算量
    cv.CHAIN_APPROX_NONE: 这种方法保留轮廓上的所有点. 即使是一个长直线段, 也会保留所有的像素点.
    速度: 这种方法速度较慢,因为它没有进行压缩,返回的轮廓点更多,计算量大.
    """
    # 返回值可能有多个,根据opencv的版本有关系,里面有cnts,还有hierarchy(轮廓关系)
    cnts = cv.findContours(board, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt)]
    if len(largeContours) >= 2:
        centerPoints = [get_center_by_contour(contour) for contour in largeContours]
        for index, center in enumerate(centerPoints):
            anotherCenter = centerPoints[(index + 1) % len(centerPoints)]
            cv.line(board, center, anotherCenter, colorWhite, 20)
        show_image("CenterLineJointed", board)
        cnts = cv.findContours(board, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
        cnts = imutils.grab_contours(cnts)
        cntFirstFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]
    elif len(largeContours) == 1:
        cntFirstFinded = largeContours[0]
    else:
        result.info = "没有找到任何轮廓"
        result.data = CheckData(None, None, None, None, None)
        return result
    # 将四个角的影响去除掉
    rectPoints = np.intp(cv.boxPoints(cv.minAreaRect(cntFirstFinded)))
    leftTop, rightTop, rightBottom, leftBottom = get_rect_points_clockwise(rectPoints)
    boardFinal = board.copy()
    moveWidth = int(get_two_point_distance(leftTop, rightTop) * 0.1)
    moveHeight = int(get_two_point_distance(leftTop, leftBottom) * 0.2)
    extend = 30
    # 左上角填黑
    x, y = leftTop
    boardFinal[y - extend:y + moveHeight, x - extend:x + moveWidth] = 0
    # 右上角填黑
    x, y = rightTop
    boardFinal[y - extend:y + moveHeight, x - moveWidth:x + extend] = 0
    # 右下角填黑
    x, y = rightBottom
    boardFinal[y - moveHeight:y + extend, x - moveWidth:x + extend] = 0
    # 左下角填黑
    x, y = leftBottom
    boardFinal[y - moveHeight:y + extend, x - extend:x + moveWidth] = 0
    # 将boardFinal做闭操作还原
    boardFinal = cv.morphologyEx(boardFinal, cv.MORPH_CLOSE, kernel, iterations=morphTimes)
    show_image("MoveAngleAndClose", boardFinal)
    cnts = cv.findContours(boardFinal.copy(), cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt) > filterArea]
    if len(largeContours) >= 1:
        cntFinalFinded = sorted(largeContours, key=cv.contourArea, reverse=True)[0]
        box = cv.minAreaRect(cntFinalFinded)
        (cx, cy), (w, h), angle = box
        angleNew = None
        if angle > 70:
            angleNew = 90 - angle
            w, h = h, w
            box = (cx, cy), (w, h), angleNew

        if angle < -70:
            angleNew = 90 + angle
            w, h = h, w
            box = (cx, cy), (w, h), angleNew
        print(f"angleFirst: {angle},angleNew:{angleNew}")

        # 获取传感器的触发位置和第一边的距离,这个怎么获取呢,传感器的位置有(x和y坐标)
        triggerLineK, triggerLineB = get_line_k_and_b(leftTop, rightTop)
        sensorDis = get_point_to_line_distance(sensorPoint, triggerLineK, triggerLineB)
        sensorDisMm = round(sensorDis // mm2pix, 2)
        # 将检测到的传感器的距离画到图上去
        cntDraw = np.array([leftTop, rightTop, rightBottom, leftBottom])
        cv.drawContours(imageColor, [cntDraw], -1, colorGreen, thickness=4)
        cv.circle(imageColor, sensorPoint, 5, colorRed, thickness=4)
        widthMm, heightMm = w // mm2pix, h // mm2pix
        topMid = get_two_points_middle(leftTop, rightTop)
        leftMid = get_two_points_middle(leftTop, leftBottom)
        heightPos = int(leftMid[0] + 20), int(leftMid[1])
        cv.putText(imageColor, f"H: {heightMm}", heightPos, cv.FONT_HERSHEY_SIMPLEX,
                   fontScale=3, color=colorBlue, thickness=2)
        sensorDisPos = int(topMid[0] - 500), int(topMid[1] + 200)
        cv.putText(imageColor, f"W: {widthMm}, SensorDis: {sensorDisMm}", sensorDisPos,
                   cv.FONT_HERSHEY_SIMPLEX, fontScale=3, color=colorBlue, thickness=2)

        show_image("sensorDis", imageColor)
        boardResNoAngle, _ = get_board_without_angle(imageColor, box)
        show_image("totalBoardResNoAngle", boardResNoAngle)
        # 这里去除掉轮廓周围的噪声点,全部弄成黑色,方便后面进行图像处理
        imageCntMask = cv.drawContours(image.copy(), [cntDraw], -1, colorWhite, -1)
        imageFillBlack = cv.bitwise_and(image, imageCntMask)
        show_image("ImageFillBlack", imageFillBlack)
        totalboard, offset = get_board_without_angle(imageFillBlack, box)
        show_image("totalboard", totalboard)
        result.data = CheckData(box, offset, imageColor, boardResNoAngle, totalboard)
        if abs(sensorDisMm) > sensorDisThres:
            result.info = (f"传感器触发异常: 物料检测第一边距离传感器超过阈值,检测距离:{sensorDisMm},"
                           f"阈值:{sensorDisThres}")
            return result
        heightMin = 1300 - 25
        heightMax = 1300 + 25
        if heightMm < heightMin or heightMm > heightMax:
            result.info(f"物料总长度检测异常,检测值:{heightMm},最小长度:{heightMin},"
                        f"最大长度:{heightMax},请确认来料是否正确")
            return result
        result.state = True
        result.info = f"物料检测通过: 检测物料总长: {heightMm},传感器触发位置偏差: {sensorDisMm}"
        return result

    else:
        result.info = "去除掉四个角之后没有找到任何轮廓"
        result.data = CheckData(None, None, None, None, None)
        return result


def get_sub_images(image):
    """
    获取图像的第一边识别的ROI,中缝识别的ROI,底边识别的ROI.以及他们和TotalBoard的Offset.
    策略就是先转移到TotalBoard,然后再转移到大图上去.
    :param image:
    :param box:
    :return:
    """
    topSlice = int(100 * mm2pix)
    midSlice = int(200 * mm2pix)
    downSlice = int(100 * mm2pix)
    subTop = image[:topSlice, :]
    topOffset = 0
    midY = image.shape[0] // 2
    subMid = image[midY - midSlice:midY + midSlice, :]
    midOffset = midY - midSlice
    subBottom = image[image.shape[0] - downSlice:image.shape[0], :]
    downOffset = image.shape[0] - downSlice
    show_image("SubTop", subTop)
    show_image("SubMid", subMid)
    show_image("SubBottom", subBottom)
    return (subMid, subTop, subBottom), (midOffset, topOffset, downOffset)


def get_angle_by_k(k):
    """
    根据k获取角度
    :param k:
    :return:
    """
    if k == 0:
        return 0
    if k is None:
        return 90
    return np.arctan(k) * 180 / np.pi
